package com.stefan.tree;

import java.util.*;

/*
* 哈夫曼树 - 带权重的树
* 可以得到哈夫曼编码， 用于数据压缩和加密等功能
* */
public class HafumanTree<E> {

    TreeNode<E> root;//树根
    List<TreeNode<E>> leaves; //data节点(也是所有的叶子节点) - 最后需要编码的，所以要存起来
    Map<E,Integer> map;//节点和对应的权重

    //构建hafuman树的条件，是不是叶子节点和他们的权重
    public HafumanTree (Map<E,Integer> map) {
        this.map = map;
        leaves = new ArrayList<>();//初始化leaves 存储
    }

    private void create() {
        if (root != null) {
            return;//已经创建过了
        }
        E[] keys = (E[]) map.keySet().toArray(); //这个就是节点里面的data
        //这里需要借助排序工具来存储权重，取2个权重最小的构建子树，用PriorityQueue，
        //这里也就是为啥要实现排序接口，让最小的权重在用PriorityQueue可以被取出
        PriorityQueue<TreeNode<E>> queue = new PriorityQueue<>();
        for (E key : keys) {
            TreeNode<E> treeNode = new TreeNode(key,map.get(key));
            queue.offer(treeNode); //
            leaves.add(treeNode); //for
        }
        //构建树，每次取2个最小的构建
        int size = queue.size();
        for (int i = 0; i < size-1; i++) {
            TreeNode<E> q1 = queue.poll();
            TreeNode<E> q2 = queue.poll();
            int var1 = q1.frequency + q2.frequency;
            TreeNode<E> treeNode = new TreeNode<>(null,var1);
            treeNode.left = q1;
            treeNode.right = q2;
            q1.parent = treeNode;
            q2.parent = treeNode;
            queue.add(treeNode);
        }
        root = queue.poll();
    }
    //对所有叶子节点编码
    public Map<E,String> hafumanCode() {
        Map<E,String> result = new HashMap<>();
        for (TreeNode<E> leaf : leaves) {
            StringBuilder br = new StringBuilder();
            TreeNode<E> current = leaf;
            while (current.parent != null) {
                if (current == current.parent.left){ //当前点是左边
                    br.append("0");
                }
                if (current == current.parent.right) {//当前点在右边
                    br.append("1");
                }
                current = current.parent;
            }
            result.put(leaf.data,br.reverse().toString());
            System.out.println(leaf.data+" : " + result.get(leaf.data));
        }
        return result;
    }
    public void encrypt(String text) {
        //创建
        create();
        Map<String, String> book = (Map<String, String>)hafumanCode(); //密码本
        char[] charArray = text.toCharArray();
        StringBuilder encrypt = new StringBuilder();
        for (char c : charArray) {
            encrypt.append(book.get(String.valueOf(c)));
        }
        System.out.println("encrypt text: "+ encrypt.toString());
    }
    public void decrypt(String text) {
        //创建
        create();
//        Map<String, String> book = (Map<String, String>)hafumanCode(); //密码本
        StringBuilder decrypt = new StringBuilder();
        char[] charArray = text.toCharArray();
        TreeNode<E> current = root;
        for (char c : charArray) {
            current = c == '0' ? current.left : current.right; //0往左，1往右，知道找到叶子节点 //所有编码元素都在叶子节点
            //到达叶子节点
            if (current.left == null && current.right == null) {
                decrypt.append(current.data);
                current = root; //每次解码都要从root节点
            }
        }
        System.out.println("decrypt text: "+ decrypt.toString());
    }


    static class TreeNode<E> implements Comparable<TreeNode<E>> {
        E data; //节点里面的字面量
        int frequency; //权重
        TreeNode<E> left;
        TreeNode<E> right;
        TreeNode<E> parent; //可以从叶子节点往上查询

        public TreeNode (E data,int frequency) {
            this.data = data;
            this.frequency = frequency;
        }

        @Override
        public int compareTo(TreeNode<E> o) {
            return this.frequency - o.frequency; //用权重排序
        }
    }

    public static void main(String[] args) {
        Map<String,Integer> map = new HashMap<>();
        map.put("a",3);
        map.put("b",24);
        map.put("c",6);
        map.put("d",20);
        map.put("e",34);
        map.put("f",4);
        map.put("g",12);

        HafumanTree<String> hafuma = new HafumanTree<>(map);
        Scanner scanner = new Scanner(System.in);
        String input = scanner.nextLine();
//        hafuma.encrypt(input);
        hafuma.decrypt(input);
    }

}
